#coding=utf-8

'''
思考：如何自制函数，返还任意概率分布的离散分布，例如泊松分布：$k = 0,1,2,\cdots$ 相应***概率***为 $f(k; \lambda)=\frac{\lambda^k e^{-\lambda}}{k!}$。
如有可能，请作图验证。
'''
import math

import numpy as np
import matplotlib.pyplot as plt

# python 中 lambda 有特殊的意义，所以这里把变量名叫做　lamb (小羊羔)
def f(k, lamb):
    return np.power(lamb, k) * np.exp(-lamb) / np.math.factorial(k)

def plotf(lamb, xmax):
    x = np.arange(0, xmax, 1)
    y = [ f(k, lamb) for k in x ]
    plt.plot( x, y )

# 理论上可以推断： k+1 > lamb, 就有　f(k+1, lamb ) < f(k, lamb)，
# 所以，如果取 k = [lamb], 则有 f(k+1, lamb)/f(k,lamb) = lamb / ([lamb]+1) < 1,
#                                f(k, lamb)/f(k-1, lamb) = lamb / [lamb] >= 1
# 所以，可以取　fmax = f( [lamb], lamb )，即 poisson 分布的一个上确界

# generate n intergs in [0, xmax), and sift them, so that the remains obey poisson distribution
def poisson( lamb, n, xmax, fmax ):
    x = np.random.randint( 0, xmax, n )
    a = np.random.random( n )
    #print("x=", x, " a=", a)
    return [ x[i] for i in range(n) if a[i] < f(x[i],lamb)/fmax ]

lamb = 2; fmax = f(math.floor(lamb), lamb); print("fmax = ", fmax)
samples = poisson( lamb, 1000, 10, fmax ); #print("samples = ", samples)
plt.hist( samples, bins=np.arange(-0.5, 10, 1 ), rwidth=0.1, density=True )
plotf(lamb, 10)
plt.show()